#pragma once

/*
 * anet thread pool with round-robin mode by asio.
 */

#include <memory>
#include <vector>
#include <thread>
#include <assert.h>
#include <functional>
#include "asio.hpp"
#include "asio/detail/noncopyable.hpp"

namespace anet {
	class CThreadPool final: asio::noncopyable {
	public:
		explicit CThreadPool(size_t poolSize) : m_nextIndex(0) {
			auto realPoolSize = poolSize;
			if (realPoolSize <= 0) {
				// !!std::thread::hardware_concurrency() may return 0
				realPoolSize = std::thread::hardware_concurrency();
			}

			// double check the default value.
			if (realPoolSize <= 0) {
				realPoolSize = 1;
			}

			// create IO context and its work.
			for (size_t i = 0; i < realPoolSize; ++i) {
				auto io_service = std::make_shared<asio::io_context>();
				auto work = std::make_shared<asio::io_context::work>(*io_service);
				m_vecIOContexts.emplace_back(io_service);
				m_vecWorks.emplace_back(work);
			}
		}
		virtual ~CThreadPool() = default;
		CThreadPool(const CThreadPool& rhs) = delete;
		CThreadPool& operator=(const CThreadPool &rhs) = delete;

		using IOContextPtr = std::shared_ptr<asio::io_context>;
		using WorkPtr = std::shared_ptr<asio::io_context::work>;
		using ThreadPtr = std::shared_ptr<std::thread>;

	public:
		// Add task to rand IO service to run.
		template<typename F, typename... Args>
		void AddTask(F &&f, Args&&... args) {
			GetIoContext().post(std::bind(f, std::forward<Args>(args)...));
		}

		// wait all threads to exit.
		void Join() {
			for (size_t i = 0; i < m_vecIOContexts.size(); i++) {
				m_vecThreads[i]->join();
			}
		}

		void Stop() {
			m_vecWorks.clear();

			for (size_t i = 0; i < m_vecThreads.size(); i++) {
				m_vecThreads[i]->join();
			}
			m_vecThreads.clear();

			for (size_t i = 0; i < m_vecIOContexts.size(); i++) {
				m_vecIOContexts[i]->stop();
			}
		}

		// start thread pool to run IO service event.
		void Init() {
			for (size_t i = 0; i < m_vecIOContexts.size(); i++) {
				auto ios = m_vecIOContexts[i];
				auto thr = std::make_shared<std::thread>(
					// start thread function to deal IO service run.
					 [this, ios]() {
					   ios->run();
				     });
				m_vecThreads.emplace_back(thr);
			}
		}

		// Get an IO context object.
		asio::io_context& GetIoContext() {
			asio::io_context& ioContext = *m_vecIOContexts[m_nextIndex];
			if (++m_nextIndex == int(m_vecIOContexts.size())) {
				m_nextIndex = 0;
			}
			return ioContext;
		}

		// last one for acceptor context.
		asio::io_context& GetAcceptorIoContext() {
			return *m_vecIOContexts.back();
		}

	protected:
		std::vector<IOContextPtr> m_vecIOContexts;
		std::vector<WorkPtr>      m_vecWorks;
		std::vector<ThreadPtr>    m_vecThreads;
		std::atomic_int           m_nextIndex;
	};
}
